import random
import numpy as np
import torch
from env import create_env
from agent import DQNAgent
from trainer import train_agent, test_agent
from config import config


def main():
    random.seed(config.SEED)
    np.random.seed(config.SEED)
    torch.manual_seed(config.SEED)

    train_env = create_env('train')
    test_env = create_env('test')

    state_dim = train_env.observation_space.shape[0]
    action_dim = train_env.action_space.n

    agent = DQNAgent(state_dim, action_dim)

    train_agent(agent, train_env)

    test_agent(agent, test_env)

    train_env.close()
    test_env.close()


if __name__ == "__main__":
    main()
